#!/usr/bin/env bash

# RabbitMQ Benchmarker
# Copyright 2021 VMware, Inc.
#
# SPDX-License-Identifier: Apache-2.0
#
# This product is licensed to you under the Apache 2.0 license
# (the "License").  You may not use this product except in compliance
# with the Apache 2.0 License.
#
# This product may include a number of subcomponents with separate
# copyright notices and license terms. Your use of these subcomponents
# is subject to the terms and conditions of the subcomponent's license,
# as noted in the LICENSE file.

set -euo pipefail

GREEN='\033[0;32m'
ORANGE='\033[0;33m'
RED='\033[0;31m'
NO_COLOR='\033[0m'

usage() {
    usage=$(
        cat <<-END
USAGE:
  Provision infrastructure and run benchmark on Kubernetes
    % benchmark --provider (gke|aks|eks|calatrava)
  Install and run Benchmarking tools on existing targeted Kubernetes
    % benchmark --skip-terraform
  Destroy infrastructure
    % benchmark destroy --provider (gke|aks|eks|calatrava)
END
    )
    echo "$usage"
}

terraform_env() {
    case "$provider" in
        "gke")
            terraform_gke
            ;;

        "aks")
            terraform_aks
            ;;

        "eks")
            terraform_eks
            ;;

        "calatrava")
            terraform_calatrava
            ;;

        *)
            printf "%sOption '%s' not recognized%s\n" "$RED" "$provider" "$NO_COLOR"
            exit 1
            ;;
    esac
}

terraform_gke() {
    pushd terraform/gke
        printf "%bTerraforming a GKE cluster.%b\n" "$GREEN" "$NO_COLOR"
        terraform init
        terraform apply -var-file="../../gke-cluster-config.tfvars"
        export KUBECONFIG="$PWD/kubeconfig-rabbitmq-benchmark"
        printf "%bCluster created, credentials are located in 'terraform/gke/kubeconfig-rabbitmq-benchmark'.%b\n" "$GREEN" "$NO_COLOR"
    popd
}

terraform_aks() {
    pushd terraform/aks
        printf "%bTerraforming an AKS cluster.%b\n" "$GREEN" "$NO_COLOR"
        terraform init
        terraform apply -var-file="../../aks-cluster-config.tfvars"
        terraform output kube_config > kubeconfig-rabbitmq-benchmark
        export KUBECONFIG="$PWD/kubeconfig-rabbitmq-benchmark"
        printf "%bCluster created, credentials are located in 'terraform/aks/kubeconfig-rabbitmq-benchmark'.%b\n" "$GREEN" "$NO_COLOR"
    popd
}

terraform_eks() {
    pushd terraform/eks
        printf "%bTerraforming an EKS cluster.%b\n" "$GREEN" "$NO_COLOR"
        terraform init
        terraform apply -var-file="../../eks-cluster-config.tfvars"
        export KUBECONFIG="$PWD/kubeconfig_benchmark"
        printf "%bCluster created, credentials are located in 'terraform/eks/kubeconfig_benchmark'.%b\n" "$GREEN" "$NO_COLOR"
    popd
}

compile_calatrava_provider() {
    calatrava_source="$(mktemp -d)"
    trap 'rm -rf "$calatrava_source"' EXIT
    git clone git@gitlab.eng.vmware.com:calatrava/calatrava.git "$calatrava_source"
    pushd "$calatrava_source/terraform"
        make install
    popd
}

terraform_calatrava() {
    printf "%bTerraforming a Calatrava cluster.%b\n" "$GREEN" "$NO_COLOR"
    compile_calatrava_provider
    pushd terraform/calatrava
        terraform init
        terraform apply -var-file="../../calatrava-cluster-config.tfvars"
        export KUBECONFIG="$PWD/kubeconfig-rabbitmq-benchmark"
        printf "%bCluster created, credentials are located in 'terraform/eks/kubeconfig-rabbitmq-benchmark'.%b\n" "$GREEN" "$NO_COLOR"
    popd
    configureClusterRoleBinding
}

configureClusterRoleBinding() {
    kubectl create clusterrolebinding default-binding \
        --clusterrole=psp:vmware-system-unprivileged \
        --group=system:authenticated
}

deploy() {
    deploy_cluster_operator
    deploy_influx_db
    deploy_kube_prometheus_stack
    deploy_production_cluster
}

deploy_cluster_operator() {
    printf "%bDeploying latest RabbitMQ cluster operator.%b\n" "$GREEN" "$NO_COLOR"
    kubectl apply -f "https://github.com/rabbitmq/cluster-operator/releases/latest/download/cluster-operator.yml"
    printf "%bCluster operator deployed.%b\n" "$GREEN" "$NO_COLOR"
}

deploy_kube_prometheus_stack() {
    printf "%bDeploying Kube Prometheus Stack.%b\n" "$GREEN" "$NO_COLOR"
    promNS="prom"
    helm repo add prometheus-community https://prometheus-community.github.io/helm-charts
    helm upgrade --install prom \
        --namespace "$promNS" --create-namespace \
        -f test-infrastructure/prom-values.yml \
        prometheus-community/kube-prometheus-stack

    operator="$(mktemp -d)"
    trap 'rm -rf "$operator"' EXIT
    pushd "$operator"
        wget -c "$(curl -s https://api.github.com/repos/rabbitmq/cluster-operator/releases/latest | jq -r '.tarball_url')" -O cluster-operator.tar.gz
        tar xfz cluster-operator.tar.gz
        source_dir="$(find "$operator" -type d -name 'rabbitmq-cluster-operator-*')"
    popd

    printf "%bDeploying ServiceMonitor and PodMonitor.%b\n" "$GREEN" "$NO_COLOR"
    kubectl -n "$promNS" apply -f "$source_dir/observability/prometheus/monitors/"

    printf "%bDeploying Prometheus Rules.%b\n" "$GREEN" "$NO_COLOR"
    kubectl -n "$promNS" apply --recursive -f "$source_dir/observability/prometheus/rules/"

    printf "%bDeploying Grafana Dashboards.%b\n" "$GREEN" "$NO_COLOR"
    kubectl -n "$promNS" apply -f "$source_dir/observability/grafana/dashboards/"

    printf "To open Grafana run\n%% kubectl -n %s port-forward svc/prom-grafana 3000:80\nand open your browser at http://localholst:3000\nusername: admin, password: admin\n\n%bKube Prometheus Stack deployed.%b\n" "$promNS" "$GREEN" "$NO_COLOR"
}

deploy_influx_db() {
    printf "%bDeploying InfluxDB for storing metrics.%b\n" "$GREEN" "$NO_COLOR"
    helm repo add influxdata https://helm.influxdata.com/
    helm upgrade --install benchmark-data \
        -f test-infrastructure/influx-values.yml \
        influxdata/influxdb
    printf "%bInfluxDB deployed.%b\n" "$GREEN" "$NO_COLOR"
}

deploy_production_cluster() {
    printf "%bDeploying RabbitMQ cluster.%b\n" "$GREEN" "$NO_COLOR"
    ytt -f test-infrastructure/rabbitmq.yml -f RabbitMQ-values.yml | kubectl apply -f-
    printf "Waiting for RabbitMQ cluster to be ready."
    while [ "$(kubectl get rmq benchmark -o jsonpath='{.status.conditions[?(@.type=="AllReplicasReady")].status}')" != "True" ]; do
        printf "."
        sleep 5
    done
    printf "\n"
    printf "%bRabbitMQ cluster deployed.%b\n" "$GREEN" "$NO_COLOR"
}

# Now to actually run the benchmark
# The idea is to use RabbitTestTool, which requires some policy and topology files
# Additionally, we should collect the test results in an InfluxDB
benchmark() {
    printf "%bRunning Rabbit Test Tool. Test results exported to InfluxDB.%b\n" "$GREEN" "$NO_COLOR"
    ytt -f test-infrastructure/benchmarker.yml -f test-infrastructure/config.json -f policy.json -f topology.json | kubectl apply -f-
}

destroy() {
    case "$provider" in
        "gke")
            destroy_gke
            return
            ;;

        "aks")
            destroy_aks
            return
            ;;

        "eks")
            destroy_eks
            return
            ;;

        "calatrava")
            destroy_calatrava
            return
            ;;

        *)
            printf "%bOption '$provider' not recognized%b\n" "$RED" "$NO_COLOR"
            exit 1
            ;;
    esac
}

destroy_gke() {
    pushd terraform/gke
        printf "%bDestroying GKE cluster.%b\n" "$ORANGE" "$NO_COLOR"
        terraform destroy
        printf "%bGKE cluster destroyed.%b\n" "$GREEN" "$NO_COLOR"
    popd
}

destroy_aks() {
    pushd terraform/aks
        printf "%bDestroying AKS cluster.%b\n" "$ORANGE" "$NO_COLOR"
        terraform destroy
        printf "%bAKS cluster destroyed.%b\n" "$GREEN" "$NO_COLOR"
    popd
}

destroy_eks() {
    pushd terraform/eks
        printf "%bDestroying EKS cluster.%b\n" "$ORANGE" "$NO_COLOR"
        terraform destroy
        printf "%bEKS cluster destroyed.%b\n" "$GREEN" "$NO_COLOR"
    popd
}

destroy_calatrava() {
    pushd terraform/calatrava
        printf "%bDestroying Calatrava cluster.%b\n" "$ORANGE" "$NO_COLOR"
        terraform destroy
        printf "%bCalatrava cluster destroyed.%b\n" "$GREEN" "$NO_COLOR"
    popd
}

case "$#" in
    1)
        if [[ "$1" == "--help" ]]; then
            usage
            exit 0
        elif [[ "$1" == "--skip-terraform" ]]; then
            deploy
            benchmark
            exit 0
        else
            usage
            exit 1
        fi
        ;;

    2)
        if [[ "$1" != "--provider" ]]; then
            usage
            exit 1
        fi
        provider="$2"
        terraform_env
        deploy
        benchmark
        exit 0
        ;;

    3)
        if [[ "$1" != "destroy" ]] || [[ "$2" != "--provider" ]]; then
            usage
            exit 1
        fi
        provider="$3"
        destroy
        exit 0
        ;;

    *)
        usage
        exit 1
        ;;
esac
